"""
简化的MCP与大模型集成演示脚本
演示大模型如何调用MCP服务的核心概念
"""

import json
from typing import Dict, Any, List


class SimpleMCPService:
    """简化的MCP服务模拟类"""
    
    def __init__(self):
        self.servers = {
            "search-server": {
                "enabled": True,
                "url": "http://localhost:3000/search",
                "tools": ["web_search", "news_search", "image_search"]
            },
            "calculator-server": {
                "enabled": True,
                "url": "http://localhost:3001/calc",
                "tools": ["add", "subtract", "multiply", "divide"]
            }
        }
    
    def is_server_available(self, server_name: str) -> bool:
        """检查服务器是否可用"""
        return server_name in self.servers and self.servers[server_name]["enabled"]
    
    def list_tools(self, server_name: str) -> List[str]:
        """列出服务器上的工具"""
        if self.is_server_available(server_name):
            return self.servers[server_name]["tools"]
        return []
    
    def call_tool(self, server_name: str, tool_name: str, parameters: Dict[str, Any]) -> Dict[str, Any]:
        """调用工具"""
        if not self.is_server_available(server_name):
            raise Exception(f"服务器 {server_name} 不可用")
        
        if tool_name not in self.servers[server_name]["tools"]:
            raise Exception(f"工具 {tool_name} 在服务器 {server_name} 上不可用")
        
        # 模拟工具调用结果
        if server_name == "search-server" and tool_name == "web_search":
            query = parameters.get("query", "")
            return {
                "results": [
                    f"关于'{query}'的搜索结果1",
                    f"关于'{query}'的搜索结果2",
                    f"关于'{query}'的搜索结果3"
                ],
                "total_results": 3
            }
        elif server_name == "calculator-server":
            a = parameters.get("a", 0)
            b = parameters.get("b", 0)
            if tool_name == "add":
                return {"result": a + b}
            elif tool_name == "subtract":
                return {"result": a - b}
            elif tool_name == "multiply":
                return {"result": a * b}
            elif tool_name == "divide":
                if b == 0:
                    raise Exception("除数不能为零")
                return {"result": a / b}
        
        return {"result": "模拟工具调用成功"}


# 创建MCP服务实例
mcp_service = SimpleMCPService()


def demonstrate_mcp_integration():
    """演示MCP集成"""
    print("=== MCP与大模型集成概念演示 ===")
    print("展示大模型如何调用MCP服务的核心概念")
    print()
    
    # 1. 列出所有MCP服务器
    print("1. 列出所有MCP服务器:")
    servers_info = {
        "enabled": True,
        "servers": list(mcp_service.servers.keys()),
        "server_configs": mcp_service.servers
    }
    print(json.dumps(servers_info, indent=2, ensure_ascii=False))
    print()
    
    # 2. 检查特定服务器可用性
    print("2. 检查search-server可用性:")
    is_available = mcp_service.is_server_available("search-server")
    availability_info = {
        "server": "search-server",
        "available": is_available
    }
    print(json.dumps(availability_info, indent=2, ensure_ascii=False))
    print()
    
    # 3. 列出服务器上的工具
    print("3. 列出search-server上的工具:")
    tools = mcp_service.list_tools("search-server")
    tools_info = {
        "server": "search-server",
        "tools": tools
    }
    print(json.dumps(tools_info, indent=2, ensure_ascii=False))
    print()
    
    # 4. 调用搜索工具
    print("4. 调用search-server的web_search工具:")
    try:
        search_result = mcp_service.call_tool(
            "search-server", 
            "web_search", 
            {"query": "最新的AI技术发展"}
        )
        print(json.dumps(search_result, indent=2, ensure_ascii=False))
    except Exception as e:
        print(f"调用工具时出错: {e}")
    print()
    
    # 5. 调用计算器工具
    print("5. 调用calculator-server的add工具:")
    try:
        calc_result = mcp_service.call_tool(
            "calculator-server", 
            "add", 
            {"a": 15, "b": 25}
        )
        print(json.dumps(calc_result, indent=2, ensure_ascii=False))
    except Exception as e:
        print(f"调用工具时出错: {e}")
    print()
    
    # 6. 演示错误处理
    print("6. 演示错误处理 - 调用不存在的工具:")
    try:
        error_result = mcp_service.call_tool(
            "search-server", 
            "nonexistent_tool", 
            {"param": "value"}
        )
        print(json.dumps(error_result, indent=2, ensure_ascii=False))
    except Exception as e:
        print(f"捕获到预期的错误: {e}")
    print()
    
    print("=== 演示结束 ===")
    print()
    print("在实际应用中，这些功能会被封装成LangChain工具，")
    print("供大模型在LangGraph工作流中调用。")
    print()
    print("核心概念包括：")
    print("1. 服务发现 - 列出可用的MCP服务器")
    print("2. 可用性检查 - 确认服务器是否在线")
    print("3. 工具枚举 - 获取服务器支持的工具列表")
    print("4. 工具调用 - 执行具体的功能")
    print("5. 错误处理 - 优雅地处理各种异常情况")


if __name__ == "__main__":
    demonstrate_mcp_integration()